import sys
sys.path.append("../../")

import tushare as ts
import numpy as np
import pandas as pd
from pandas import Series, DataFrame
from frameworks.cores.Config import *

class TushareCalc:
    pro = ''
    df = ''
    code = ''
    start_date = ''
    end_date = ''
    def __init__(self,code="000001.SZ",start_date='',end_date=''):
        token = Config().getTuShareKey()
        ts.set_token(token)
        # 初始化pro接口
        self.pro = ts.pro_api()
        self.code = code
        self.start_date = start_date
        self.end_date = end_date

    def getCodeInfo(self):
        data = self.pro.query('stock_basic', exchange='', list_status='L',fields='ts_code,symbol,name,industry,list_date')
        for i, row1 in enumerate(data.values):
            if(self.code == row1[1]):
                return row1

    def queryCode(self):
        self.df = self.pro.daily(ts_code=self.code, start_date=self.start_date, end_date=self.end_date)
        self.df.to_csv('a1.csv')  # 这里随意取名，这个csv会存在当前.py文件的同一目录下，可以查看具体数据
        self.df = pd.read_csv('a1.csv', parse_dates=[
            'trade_date'])  # 后面这个parse_dates参数的作用是将trade_date那一列的日期数据从xxxxxxxx（float型）改成xxxx-xx-xx（日期型）形式
        self.df.drop(labels='Unnamed: 0', axis=1, inplace=True)
        self.df = self.df.sort_values(by='trade_date')
        """
        这一步是让我们的数据以日期从小到大排序，因为之前我们获取数据的时候print出来发现最前面的是最近的数据，而我们画均线需要按时间发展来画，所以应当以日期排序
        """
        self.df.set_index('trade_date', inplace=True)  # 把索引设为日期
        # print(df)#这里也可以print一下来确认自己的操作对不对


    def getMA(self,n):
        self.queryCode()
        return self.df['close'].rolling(n).mean()

    def getClose(self):
        self.queryCode()
        return self.df['close']
